#!/usr/bin/env python
#-*- coding: utf-8 -*-

# Author: Edward Roualdes; edited code from Tommy Bullock and Rachel Schwartz
#         University of Kentucky

import math
import os
import re
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
import utils
from sanity import sane

# some variables
bases = ['A', 'C', 'G', 'T', 'a', 'c', 'g', 't', '-']
imperfect_links = {}          #dictionary containing locus_indiv: # imperfect linkages

# functions
def l(x):
    "natural log function called faster than using math member fn."
    return math.log(x)

def likelihoodtest(list_bp):
    """calculate the likelihood ratio test of Hohenlohe 2010 - takes a sequence in list form."""
    
    # dictionary of base frequencies
    bpfreq={'A':0, 'C':0, 'G':0, 'T':0, 'a':0, 'c':0, 'g':0, 't':0, '-':0} 
    bases=bpfreq.keys()         #get list of bases

    for base in bases:  
        bpfreq[base]=list_bp.count(base) # edit dictionary of base frequencies by counting the number of each base in the sequence

    all=sorted(bpfreq.items(), key=lambda (k,v): v) # sort bases by frequency and put in list
    n1=all[-1][1]                                 # get base frequencies and associated bases in order of frequency (e.g. n1 is most frequent)
    n1=float(n1)
    base1=all[-1][0]
    n2=all[-2][1]
    n2=float(n2)
    base2=all[-2][0]
    n3=all[-3][1]
    n3=float(n3)
    n4=all[-4][1]
    n4=float(n4)
    numbp=float(n1+n2+n3+n4)

    #function from Hohenlohe 2010
    likelihood = n1*l(n1/numbp) 
    if numbp-n1 != 0:
        likelihood += (numbp - n1)*l((numbp - n1)/(3*numbp))
        
    if n1+n2 != 0:
        likelihood -= (n1 + n2)*l((n1+n2)/(2*numbp))

    if n3+n4 != 0:
        likelihood -= (n3 + n4)*l((n3+n4)/(2*numbp))

    if likelihood>=3.84:        # alpha = 0.05 for test
          if (base1 == '-'):    # mark "insertions" for later deletion
               base1='d'
          return base1, base1  # return correct bases to go into allele
    elif likelihood<=(-3.84):
          return base1, base2
    else:
          base1='n'
          return base1, base1

def find2alleles(data, indiv, locus):
    """find alleles for each individual locus combination. makes all values in data (i.e. sequences) in list form (separate bases in sequence) within dictionary "data" - e.g. sequence : 'a' 'g' 'c' 't'"""

    # make dict of k = seqname, v = seq as list
    for k, v in data.iteritems():
        data[k]=list(v)
    
    # make "gaps" at ends of aligned sequences into N's
    # distinguishes between unknown data and potential gaps
    for k, v in data.iteritems():          # for each sequence      
        l = len(v)
        for i in xrange(l):
            if v[i] == '-':
                v[i] = 'n'
            else:
                break
        for i in xrange(l-1, -1, -1):
            if v[i] == '-':
                v[i] = 'n'
            else:
                break
    
    # make 2 alleles of n's
    nums = data.keys()          # list keys for data
    if nums:
        seqlength = len(data[nums[0]])
        allele1 = ['n']*seqlength
        allele2 = ['n']*seqlength
        freq_imperfect_linkage = 0 # need to know how many snps could have both possible linkages
          
        allsnps = {}  # dictionary of "position: bases for each allele"
        lastsnps = {}
        lastsnppos = -1   
          
        # go through all positions and make into two alleles
        for i in xrange(seqlength):

            snpbases = []
            for k, v in data.iteritems(): # put all bases for that position in an array
                if v[i] in bases:
                    snpbases.append(v[i])
              
            if len(snpbases) > 0:
                snp_lrt = likelihoodtest(snpbases) # get the one or two bases that should be at this position based on the LRT
                snp1 = list(snp_lrt)
            else:
                snp1 = ['n','n']
        
            lsnp1 = len(snp1)
            if lsnp1 > 1:
                if snp1[0] != snp1[1]: # sort the two alleles in snp1 if necessary
                    if lastsnps:       # if there was a previous snp
                        currentsnps = {}
                        currentsnps[snp1[0]] = []
                        currentsnps[snp1[1]] = []
                        for k, v in data.iteritems(): # dictionary of base1: seq names; base2: seq names
                            if v[i] == snp1[0]:
                                currentsnps[snp1[0]].append(k)
                            elif v[i] == snp1[1]:
                                currentsnps[snp1[1]].append(k)

                        d = set(currentsnps[snp1[0]]).intersection(set(lastsnps[lastsnp1[0]]))
                        e = set(currentsnps[snp1[1]]).intersection(set(lastsnps[lastsnp1[1]]))
                        f = set(currentsnps[snp1[0]]).intersection(set(lastsnps[lastsnp1[1]]))
                        g = set(currentsnps[snp1[1]]).intersection(set(lastsnps[lastsnp1[0]]))
                        
                        if not d and not e and not f and not g: # if no linkage
                            
                            print("Can't link SNPs correctly for ", indiv, "locus ",locus)
                            return (allele1,allele2,freq_imperfect_linkage)

                        elif (len(f)+len(g))<(len(d)+len(e)):       #if snps are linked correctly

                            if ((len(f)+len(g))>0):

                                freq_imperfect_linkage += 1 # linkage isn't perfect - add 1 

                            # if both linkages are unlikely
                            # go with the good one
                            if (len(d)+len(e))>(3*(len(f)+len(g))): 

                                lastsnps = {}
                                lastsnps = dict(currentsnps)
                                lastsnppos = i
                                lastsnp1 = list(snp1) 
                        
                            else: # if both linkages seem possible stick with n's

                                  snp1 = ['n','n']

                        elif (len(d)+len(e))<(len(f)+len(g)): # if there are more shared if the alleles were swapped
                            if ((len(d)+len(e))>0):

                                # linkage isn't perfect - add 1 
                                freq_imperfect_linkage += 1 

                            if (len(f)+len(g))>(3*(len(d)+len(e))): # if both linkages are unlikely go with the good one  
                                x = snp1[0] # swap order of snps so they are linked correctly and go with them
                                snp1[0] = snp1[1]
                                snp1[1] = x
                                lastsnps = {}
                                lastsnps = dict(currentsnps)
                                lastsnppos = i
                                lastsnp1 = list(snp1)  

                            else:

                                snp1=['n','n'] # if both linkages seem possible stick with n's

                    else:

                        lastsnps={}
                        lastsnps[snp1[0]] = []
                        lastsnps[snp1[1]] = []

                        for k, v in data.iteritems(): # dictionary of base1: seq names; base2: seq names
                            if v[i] == snp1[0]:
                                lastsnps[snp1[0]].append(k)
                            elif v[i] == snp1[1]:
                                lastsnps[snp1[1]].append(k)

                        lastsnppos = i
                        lastsnp1 = list(snp1)    
                             
                    allele1[i] = snp1[0]             
                    allele2[i] = snp1[1] # put each base into alleles
                elif snp1[0] == snp1[1]:
                    allele1[i] = snp1[0]             
                    allele2[i] = snp1[1] # put each base into alleles
            elif lsnp1 == 1:
                allele1[i] = snp1[0]             
                allele2[i] = snp1[0] # put each base into alleles
            else:
                allele1[i] = 'n'             
                allele2[i] = 'n' # put each base into alleles
          
        while "d" in allele1:
            allele1.remove("d")
            allele2.remove("d")
              
        return (allele1, allele2, freq_imperfect_linkage)

def trim_alleles(allele1, allele2):
    """take two alleles and trim n's at each end"""
    for i in xrange(len(allele1)):
        if allele1[i] == 'n':
            if allele2[i] == 'n':
                allele1[i] = 'd'
                allele2[i] = 'd'
            else:
                break
        else:
            break
    for i in xrange(len(allele1)-1, -1, -1):
        if allele1[i] == 'n':
            if allele2[i] == 'n':
                allele1[i] = 'd'
                allele2[i] = 'd'
            else:
                break
        else:
            break
    while "d" in allele1:
        allele1.remove("d")
        allele2.remove("d")
              
    return allele1, allele2

def approx2alleles(data, indiv, locus):
    """approximate alleles for each individual locus combination. makes all values in data (i.e. sequences) in list form (separate bases in sequence) within dictionary "data" - e.g. sequence : 'Y' 'R' 'W' 'S' 'K' 'M'"""

    # make dict of k = seqname, v = seq as list
    for k, v in data.iteritems():
        data[k]=list(v)
    
    # make "gaps" at ends of aligned sequences into N's
    # distinguishes between unknown data and potential gaps
    for k, v in data.iteritems():          # for each sequence      
        l = len(v)
        for i in xrange(l):
            if v[i] == '-':
                v[i] = 'n'
            else:
                break
        for i in xrange(l-1, -1, -1):
            if v[i] == '-':
                v[i] = 'n'
            else:
                break
    
    # make 2 alleles of n's
    nums = data.keys()          # list keys for data
    if nums:
        seqlength = len(data[nums[0]])
        allele1 = ['n']*seqlength
        allele2 = ['n']*seqlength
          
        # go through all positions and make into two alleles
        for i in xrange(seqlength):

            snpbases = []
            for k, v in data.iteritems(): # put all bases for that position in an array
                if v[i] in bases:
                    snpbases.append(v[i])
              
            if len(snpbases) > 0:
                snp_lrt = likelihoodtest(snpbases) # get the one or two bases that should be at this position based on the LRT
                snp1 = list(snp_lrt)
            else:
                snp1 = ['n','n']

            if snp1[0] != snp1[1]:
                if snp1[0] in 'cC' and snp1[1] in 'tT':
                    allele1[i] = 'y'
                    allele2[i] = 'y'
                elif snp1[0] in 'aA' and snp1[1] in 'gG':
                    allele1[i] = 'r'
                    allele2[i] = 'r'
                elif snp1[0] in 'aA' and snp1[1] in 'tT':
                    allele1[i] = 'w'
                    allele2[i] = 'w'
                elif snp1[0] in 'gG' and snp1[1] in 'cC':
                    allele1[i] = 's'
                    allele2[i] = 's'
                elif snp1[0] in 'tT' and snp1[1] in 'gG':
                    allele1[i] = 'k'
                    allele2[i] = 'k'
                elif snp1[0] in 'cC' and snp1[1] in 'aA':
                    allele1[i] = 'm'
                    allele2[i] = 'm'
                else:
                    allele1[i] = 'n'             
                    allele2[i] = 'n'
            else:
                allele1[i] = snp1[0]             
                allele2[i] = snp1[1]             

        while "d" in allele1:
            allele1.remove("d")
            allele2.remove("d")
              
        return allele1, allele2


def callHaps(outdir, pat = None, approx = True):
    """phase haplotypes from aligned reads and write to file."""
    # opts = sane(configfile)     # get keyword for edited Primary Alignments
    utils.mkdir(os.path.join(outdir, 'haplotypes'),
                overwrite = True,
                opt = "Please move existing \'haplotpyes\' folder ")
    data = utils.files2dict(outdir, ".*\\.aligned$", pat)
    print("Phasing haplotypes...")
    
    for d, indiv, locus in data:
        if d.keys():

            if approx:

                # approximate alleles
                allele1, allele2 = approx2alleles(d, indiv, locus)
                allele1, allele2 = trim_alleles(allele1, allele2)

            else:

                # find alleles
                allele1, allele2, freq_imperfect_linkage = find2alleles(d, indiv, locus)
                allele1, allele2 = trim_alleles(allele1, allele2) # trim ns
                
                # put freq_imperfect_linkage in imperfect_links dictionary
                locus_indiv = locus + "_" + indiv
                imperfect_links[locus_indiv] = freq_imperfect_linkage 

                    
            # create seq_records from the alleles => use seqIO.write
            a1 = SeqRecord(Seq("".join(allele1)))
            a1.id = indiv + "_1"
            a1.description = ""
        
            a2 = SeqRecord(Seq("".join(allele2)))
            a2.id = indiv + "_2"
            a2.description = ""

            # filter empty haplotypes
            alleles = filter(lambda x: x.seq.tostring() != "", [a1, a2])

            # write haplotypes to file
            if alleles:
                seqFile = open(os.path.join(outdir, 'haplotypes', locus+'.fna'), 'a')
                SeqIO.write(alleles, seqFile, "fasta")
                seqFile.close()
            del a1, a2, alleles      # del records => don't use in next loop

    if not approx:
        # frequency of imperfect linkages to file
        fileout = open(os.path.join(outdir, 'imperfect_linkages.txt'), 'a')
        for k, v in imperfect_links.iteritems():      
            fileout.write("%s\t%s\n" % (k, str(v)))
        fileout.close()

    return os.path.join(outdir, 'haplotypes')
